from logging import getLogger
from typing import BinaryIO, Optional, Union
import aioboto3
from fastapi import UploadFile
from app.core.storage.base import BaseStorage

logger = getLogger(__name__)


class S3Storage(BaseStorage):
    """S3存储实现"""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.session = aioboto3.Session(
            aws_access_key_id=self.access_key, aws_secret_access_key=self.secret_key, region_name=self.region
        )

    async def _get_client(self):
        """获取S3客户端"""
        return self.session.client("s3", endpoint_url=self.endpoint, region_name=self.region)

    async def save(self, file: Union[BinaryIO, UploadFile], filename: str) -> str:
        """保存文件到S3"""
        try:
            # 生成存储路径
            storage_path = self.generate_storage_path(filename)

            async with await self._get_client() as client:
                # 上传文件
                if isinstance(file, UploadFile):
                    await client.upload_fileobj(
                        file.file, self.bucket, storage_path, ExtraArgs={"ContentType": self.get_content_type(filename)}
                    )
                else:
                    await client.upload_fileobj(
                        file, self.bucket, storage_path, ExtraArgs={"ContentType": self.get_content_type(filename)}
                    )

            return storage_path

        except Exception as e:
            logger.error(f"S3上传失败: {str(e)}")
            raise Exception("文件上传失败")

    async def delete(self, file_path: str) -> bool:
        """从S3删除文件"""
        try:
            async with await self._get_client() as client:
                await client.delete_object(Bucket=self.bucket, Key=file_path)
            return True
        except Exception as e:
            logger.error(f"S3删除失败: {str(e)}")
            return False

    async def get(self, file_path: str) -> Optional[BinaryIO]:
        """从S3获取文件"""
        try:
            async with await self._get_client() as client:
                response = await client.get_object(Bucket=self.bucket, Key=file_path)
                return response["Body"]
        except Exception as e:
            logger.error(f"S3获取文件失败: {str(e)}")
            return None

    async def get_url(self, file_path: str, expires_in: int = 3600) -> str:
        """获取S3文件URL"""
        try:
            if self.base_url:
                return f"{self.base_url}/{file_path}"

            async with await self._get_client() as client:
                url = await client.generate_presigned_url(
                    "get_object", Params={"Bucket": self.bucket, "Key": file_path}, ExpiresIn=expires_in
                )
                return url
        except Exception as e:
            logger.error(f"生成URL失败: {str(e)}")
            raise Exception("生成文件访问URL失败")
